module hadley_merid_test_mod

  !Kent et al 2014, Dynamical core model intercomparison project: Tracer transport test cases.
  ! DCMIP1-2

	use const_mod
	use namelist_mod
	use sphere_geometry_mod
	use mesh_mod
	use vert_coord_mod
	use state_mod
	use tracer_mod
	use parallel_mod

	implicit none

	private
	public hadley_merid_test_init
	public hadley_merid_test_set_ic
	public hadley_merid_test_set_wind

	real(r8), parameter :: period = 86400     ! period of motion (1 days)
	real(r8), parameter :: u0     = 40				!	reference zonal velocity
	real(r8), parameter :: w0     = 0.15      ! reference vertical velocity
	real(r8), parameter :: k0     = 5         ! number of overturning cells   
	real(r8), parameter :: z1     = 2000			! initial longitude of first tracer
	real(r8), parameter :: z2     = 5000			! initial longitude of second tracer
	real(r8), parameter :: T0     = 300       ! isothermal atmospheric temperature
	real(r8)  ztop, p0, rho0
contains

	subroutine hadley_merid_test_init()

		ztop = 12000._r8
		p0   = 100000._r8
		rho0 = p0 / (Rd * T0)

	end subroutine hadley_merid_test_init

	subroutine hadley_merid_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, k, itracer
		real(r8) p, z, z0

		if (ntracers /= 1) then
			print*, 'ntracers must be 1!'
			stop
		end if

    z0 = 0.5 * (z1 + z2)
    do k = mesh%full_lev_ibeg, mesh%full_lev_iend
    	p = hybrid_coord_calc_ph(k, p0)
    	z = Rd * T0 / g * log(p0 / p)
    	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
    		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
    			if (z > z1 .and. z < z2) then
    				tracer%q(i,j,k,1) = 0.5 * (1 + cos(pi2 * (z - z0) / (z2 - z1)))
    			else 
    				tracer%q(i,j,k,1) = 0
    			end if
    		end do
    	end do
    end do

    do itracer = 1, size(tracer%q, 4)
    	call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    	call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
    end do

	end subroutine hadley_merid_test_set_ic

	subroutine hadley_merid_test_set_wind(mesh, state, time_in_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: time_in_seconds
		real(r8) lon, lat, p, ptop, cos_lat, cos_t, z, rho
		integer i, j, k

		ptop = p0 * exp(-g * ztop / Rd / T0)
		cos_t = cos(pi * time_in_seconds / period)

		associate (u => state%u, v => state%v, we => state%we, phs => state%phs)

		phs = p0

		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%full_lat_ibeg+1, mesh%full_lat_iend-1
				cos_lat = mesh%full_cos_lat(j)
				do i = mesh%half_lon_ibeg, mesh%half_lon_iend
					p = hybrid_coord_calc_ph(k, phs(i,j))
					z = Rd * T0 / g * log(p0 / p)
					u(i,j,k) = u0 * cos_lat
				end do
			end do
		end do

		do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			do j = mesh%half_lat_ibeg, mesh%half_lat_iend
				lat = mesh%half_lat(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					p = hybrid_coord_calc_ph(k, phs(i,j))
					z = Rd * T0 / g * log(p0 / p)
					rho = p / (Rd * T0)
					v(i,j,k) = -radius * w0 * pi * rho0 / (k0 * ztop * rho) * cos(lat) * sin(k0 * lat) * cos(pi * z / ztop) * cos_t
				end do
			end do
		end do

	  call fill_zonal_halo_lon(u, all_halo=.true.)

	  do k = mesh%half_lev_ibeg, mesh%half_lev_iend
	  	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
	  		lat = mesh%full_lat(j)
	  		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
	  			p = hybrid_coord_calc_ph_lev(k, phs(i,j))
	  			z = Rd * T0 / g * log(p0 / p)
	  			rho = p / (Rd * T0)
	  			we(i,j,k) = - g * w0 * rho0 / k0 * (-2 * sin(k0 * lat) * sin(lat) + k0 * cos(lat) * cos(k0 * lat)) * sin(pi * z / ztop) * cos_t
	  		end do
	  	end do
	  end do
	  end associate
	end subroutine hadley_merid_test_set_wind

end module hadley_merid_test_mod